2. Model Selection and Design¶
In this section we describe the model selection process and also the design of the model. We decided to use a U-Net architecture, as originally proposed in https://arxiv.org/abs/1505.04597 for biomedical image segmentation. The reason for this are threefold:
- At the original time of release, this architecture proved to be groundbreaking and set new state-of-the-art results in (biomedical) image segmentation. Although recent advancement in image segmentation models improved on these results, it can still be considered as a solid baseline model that should perform at least better than other, more general architectures;
- The design is straightforward and also the model size (in term of memory and parameters) is easy to handle (compared to current state-of-the-art architectures);
- The U-Net architecture yields a relatively good (training) performance even with relatively few sample images.
Of course there would have been many alternatives, up to transfer-learning a state-of-the-art pretrained model, such as SWIN (https://arxiv.org/abs/2103.14030) a transformer model (which are now ubiquitous in deep learning). But given the reasons above, time constraints, and the desire to code the model from the ground up and see how good it performs as a baseline / straightforward solution, we opted for the use of the U-Net architecture.
The U-Net architecture can be understood as an encoder / decoder convolutional network with the additional trait, that it allows at each encoding stage for not only a contracting but also an expanding path, where the output of the encoder layer is used as input in the corresponding decoder layer, which enables localization of features. The original architecture is visualized in the following figure.
The main differences in our implementation are, that we do not use the same input / output size and we set at each convolution a zero padding border of size 1 (i.e. we pad the borders of the input), which not only yields different input / output sizes as in the picture, but also fixes one issue the original architecture had: For each expanding path the output image needed to be cropped to fit together into the decoder layer and the output of the model would not attain the same dimensions as the input.
Otherwise, our implementation is the same, having two convolutional layers at each encoder / decoder layer and also performing max pooling at the end of each encoder layer and a transposed convolution at the end of each decoder layer (except the last one of course).
For the training we use the ADAM optimizer, which is a popular extension of stochastic gradient descent (SGD) algorithm utilizing additionally a momentum (the second moment of the gradient) and a variable learning rate for the parameters.
We use the binary cross-entropy (BCE) loss as our loss function but extend our approach at a later run to test the performance improvement when using the so called dice loss (and an average of the BCE and the dice loss). We will explain the reasoning of using the dice loss later.
We train the model for 100 epochs with a batch size of 4. We have tried different batch sizes at each model improvement iteration but they did not consistently lead to better results. Additionally we use early stopping, were we discontinue training if the validation loss does not improve over 11 epochs. Finally, we select the best model using the validation loss (and not just the model from the last epoch iteration). Let us first use just the default learning rate and see how we fare when ignoring our findings from EDA.
